#![allow(non_upper_case_globals)]
#![allow(non_camel_case_types)]
#![allow(non_snake_case)]

use std::os::raw::{c_char, c_void};

use melior::{
    Context,
    dialect::DialectRegistry,
    ir::{Module, Type, TypeLike, Value, ValueLike},
    utility::register_all_llvm_translations,
};
use mlir_sys::{
    MlirContext, MlirDialectRegistry, MlirOperation, MlirPass, MlirStringRef, MlirType, MlirValue,
};

use crate::{
    ast::{VarDeclExpr, VarType},
    lexer::Location,
};

include!(concat!(env!("OUT_DIR"), "/bindings.rs"));

pub fn mlir_load_toy_dialect(context: &Context) {
    let raw_context = context.to_raw();
    unsafe {
        mlirLoadToyDialect(raw_context);
    }
}

pub struct ScopedHashTable {
    pub ptr: *mut c_void,
}
impl ScopedHashTable {
    pub fn new() -> Self {
        Self {
            ptr: unsafe { mlirScopedHashTableCreate() },
        }
    }
}
impl Drop for ScopedHashTable {
    fn drop(&mut self) {
        unsafe { mlirScopedHashTableDestroy(self.ptr) }
    }
}

pub struct ScopedHashTableScope {
    pub ptr: *const c_void,
}
impl<'a> ScopedHashTableScope {
    pub fn new(scoped_hash_table: &ScopedHashTable) -> Self {
        Self {
            ptr: unsafe { mlirScopedHashTableScopeCreate(scoped_hash_table.ptr) },
        }
    }
}
impl Drop for ScopedHashTableScope {
    fn drop(&mut self) {
        unsafe { mlirScopedHashTableScopeDestroy(self.ptr) }
    }
}

#[repr(C)]
pub struct MlirVarDecl {
    pub name: MlirStringRef,
    pub type_name: MlirStringRef,
    pub shape_size: usize,
    pub shape: *const u64,
}
impl From<&VarDeclExpr> for MlirVarDecl {
    fn from(value: &VarDeclExpr) -> Self {
        Self {
            name: MlirStringRef {
                data: value.name.as_ptr() as *const c_char,
                length: value.name.len(),
            },
            type_name: MlirStringRef {
                data: value.r#type.name.as_ptr() as *const c_char,
                length: value.r#type.name.len(),
            },
            shape_size: value.r#type.shape.len(),
            shape: value.r#type.shape.as_ptr(),
        }
    }
}

fn mlir_string_ref_to_string(raw_string_ref: MlirStringRef) -> String {
    let bytes = unsafe {
        std::slice::from_raw_parts(raw_string_ref.data as *const u8, raw_string_ref.length)
    };
    String::from_utf8_lossy(bytes).into_owned()
}

fn mlir_shape_to_shape(shape: *const u64, shape_size: usize) -> Vec<u64> {
    let bytes = unsafe { std::slice::from_raw_parts(shape, shape_size) };
    bytes.to_vec()
}

impl Into<VarDeclExpr> for MlirVarDecl {
    fn into(self) -> VarDeclExpr {
        let name = mlir_string_ref_to_string(self.name);
        let type_name = mlir_string_ref_to_string(self.type_name);
        let shape = mlir_shape_to_shape(self.shape, self.shape_size);

        VarDeclExpr {
            name,
            r#type: VarType {
                name: type_name,
                shape,
            },
            init_val: None,
            location: Location::default(),
        }
    }
}

#[repr(C)]
pub struct MlirVarDeclValuePair {
    pub raw_value: MlirValue,
    pub raw_var: MlirVarDecl,
}

pub fn mlir_declare(scoped_hash_table: &ScopedHashTable, var: &VarDeclExpr, value: Value) -> bool {
    let raw_var_decl: MlirVarDecl = var.into();
    let raw_value = value.to_raw();

    let result = unsafe { mlirDeclare(scoped_hash_table.ptr, raw_var_decl, raw_value) };
    result == 1
}

pub fn mlir_lookup<'a>(
    scoped_hash_table: &ScopedHashTable,
    name: &str,
) -> Option<(Value<'a, 'a>, VarDeclExpr)> {
    let raw_pair: MlirVarDeclValuePair = unsafe {
        mlirLookup(
            scoped_hash_table.ptr,
            name.as_ptr() as *const c_char,
            name.len(),
        )
    };
    if raw_pair.raw_value.ptr.is_null() {
        None
    } else {
        let value = unsafe { Value::from_raw(raw_pair.raw_value) };
        let var: VarDeclExpr = raw_pair.raw_var.into();
        Some((value, var))
    }
}

#[repr(C)]
pub struct ShapeInference {
    pub ptr: *mut c_void,
}
impl ShapeInference {
    pub fn infer_shapes(&self) {
        unsafe { mlirShapeInferenceInferShapes(self.ptr) }
    }
}

pub fn mlir_dyn_cast_shape_inference(
    op_ref: melior::ir::OperationRef<'_, '_>,
) -> Option<ShapeInference> {
    let raw_op = op_ref.to_raw();
    let ptr = unsafe { mlirDynCastShapeInference(raw_op) };
    if ptr.is_null() {
        None
    } else {
        Some(ShapeInference { ptr })
    }
}

pub fn mlir_set_private(op_ref: melior::ir::OperationRef<'_, '_>) {
    let raw_op = op_ref.to_raw();
    unsafe {
        mlirSetPrivate(raw_op);
    }
}

pub fn mlir_func_register_all_extensions(registry: &DialectRegistry) {
    let raw_registry = registry.to_raw();
    unsafe {
        mlirFuncRegisterAllExtensions(raw_registry);
    }
}

pub fn mlir_llvm_register_inliner_interface(registry: &DialectRegistry) {
    let raw_registry = registry.to_raw();
    unsafe {
        mlirLLVMRegisterInlinerInterface(raw_registry);
    }
}

pub fn mlir_dump_llvm_ir<'c>(context: &'c Context, module: &Module<'c>, enable_opt: bool) -> bool {
    register_all_llvm_translations(context);
    let raw_operation = module.as_operation().to_raw();

    return unsafe { mlirDumpLLVMIR(raw_operation, enable_opt as i32) } == 0;
}

pub fn mlir_run_jit(module: &Module, enable_op: bool) -> bool {
    let raw_operation = module.as_operation().to_raw();
    return unsafe { mlirRunJit(raw_operation, enable_op as i32) } == 0;
}

pub fn mlir_struct_type_from(element_types: Vec<Type>) -> Type {
    let raw_element_types: Vec<MlirType> = element_types.into_iter().map(|t| t.to_raw()).collect();
    let raw_type =
        unsafe { mlirStructTypeFrom(raw_element_types.as_ptr(), raw_element_types.len()) };
    unsafe { melior::ir::Type::from_raw(raw_type) }
}

pub fn mlir_get_struct_element_type<'c>(struct_type: Type<'c>, index: usize) -> Option<Type<'c>> {
    let raw_struct_type = struct_type.to_raw();
    let raw_type = unsafe { mlirGetStructElementType(raw_struct_type, index) };
    if raw_type.ptr.is_null() {
        None
    } else {
        Some(unsafe { Type::from_raw(raw_type) })
    }
}
